import torch
from torch import nn
from torch.nn import init

#Efficient Channel Attention


class ECAAttention(nn.Module):

    def __init__(self,kernel_size):
        super().__init__()
        #提取每个通道的全局信息，为后续的通道注意力计算做准备。
        self.gap=nn.AdaptiveAvgPool2d(1) #全局平均池化层，将每个通道的空间维度缩小到 1x1。
        self.conv=nn.Conv1d(1,1,kernel_size=kernel_size,padding=(kernel_size-1)//2)#一维卷积层，用于跨通道的信息交互，输入和输出通道数均为 1。
        self.sigmoid=nn.Sigmoid()#Sigmoid 激活函数，将输出映射到 [0,1] 区间，得到通道注意力权重。
    #初始化权重的方法（可选）
    def init_weights(self):
        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                init.kaiming_normal_(m.weight, mode='fan_out')
                if m.bias is not None:
                    init.constant_(m.bias, 0)
                elif isinstance(m, nn.BatchNorm2d):
                    init.constant_(m.weight, 1)
                    init.constant_(m.bias, 0)
                elif isinstance(m, nn.Linear):
                    init.normal_(m.weight, std=0.001)
                    if m.bias is not None:
                        init.constant_(m.bias, 0)

    def forward(self, x):
        y = self.gap(x)  # 输出形状为 (batch_size, channels, 1, 1)
        #将通道维度放到最后，准备进行一维卷积操作，以实现通道间的信息交互。
        y = y.squeeze(-1).permute(0, 2, 1)  # 调整形状为 (batch_size, 1, channels)
        #使用一维卷积在通道维度上进行局部跨通道交互，卷积核大小为 kernel_size，输入和输出通道数均为 1。
        #通过一维卷积捕获通道间的局部依赖关系，避免了全连接层带来的高计算成本，实现高效的通道注意力机制。
        y = self.conv(y)  # 通过一维卷积，输出形状为 (batch_size, 1, channels)
        y = self.sigmoid(y)  # 经过 Sigmoid 激活，得到通道注意力权重
        #将注意力权重调整回与输入特征相兼容的形状，以便进行逐元素相乘。
        y = y.permute(0, 2, 1).unsqueeze(-1)  # 调整形状回到 (batch_size, channels, 1, 1)
        return x * y.expand_as(x)  # 将注意力权重应用到输入 x 上，进行逐通道的加权

if __name__ == '__main__':
    DEVICE = 'mps' if torch.backends.mps.is_available() else 'cpu'  # 支持 MPS 则使用，否则使用 CPU
    print(f"Using device: {DEVICE}")
    block = ECAAttention(kernel_size=3).to(DEVICE)
    input = torch.rand(1, 64, 64, 64).to(DEVICE)  # 输入张量形状为 (batch_size, channels, height, width)
    output = block(input)
    print(input.size(), output.size())